import numpy as np
import matplotlib.pyplot as plt

from pathlib import Path
from sklearn.preprocessing import OneHotEncoder
from sklearn.metrics import classification_report
from sklearn.model_selection import train_test_split
from tensorflow.keras.optimizers import SGD

from train_pac.normal.SimpleDataSetLoader import Simple_DataSet_Loader
from train_pac.normal.SimpleProcessor import Simple_Processor
from img_2_array_preprocess import Img_2_Array_Preprocess
from shallow_net import Shallow_Net


def main():
    sp = Simple_Processor(32, 32)
    i2a = Img_2_Array_Preprocess()

    sdl = Simple_DataSet_Loader(preprocessors=[sp, i2a])
    data, labels = sdl.load([str(i) for i in
                             Path('../../dataset/dog_cat').iterdir()][::5])
    data = data.astype("float") / 255.0

    # 分割数据集
    train_x, test_x, train_y, test_y = train_test_split(data, labels,
                                                        test_size=0.25,
                                                        random_state=42)
    o_hot = OneHotEncoder()
    train_y = o_hot.fit_transform(np.array([[i] for i in train_y])).toarray()
    test_y = o_hot.transform(np.array([[i] for i in test_y])).toarray()

    # 构建网络
    print("[info]:编译网络中...")
    opt = SGD(0.005)
    model = Shallow_Net.build(width=32, height=32, depth=3, classes=2)
    model.compile(optimizer=opt, loss="binary_crossentropy",
                  metrics=["acc"])

    # 训练网络
    print("[info]:开始训练网络...")
    record = model.fit(train_x, train_y, validation_data=(test_x, test_y),
                       batch_size=32, epochs=100, verbose=1)

    # 评估网络
    print("[info]:开始评估网络...")
    predictions = model.predict(test_x, batch_size=32)
    print(classification_report(test_y.argmax(1),
                                predictions.argmax(1),
                                target_names=o_hot.categories_[0]))

    # 画图
    plt.style.use("ggplot")
    plt.figure()
    plt.plot(np.arange(0, 100), record.history["loss"], label="train_loss")
    plt.plot(np.arange(0, 100), record.history["val_loss"], label="val_loss")
    plt.plot(np.arange(0, 100), record.history["acc"], label="train_acc")
    plt.plot(np.arange(0, 100), record.history["val_acc"], label="val_acc")
    plt.title("Training Loss and Accuracy")
    plt.xlabel("Epoch #")
    plt.ylabel("Loss/Accuracy")
    plt.legend()
    plt.show()

    # 模型保存
    model.save('shallowNet_animal.hdf5')


if __name__ == '__main__':
    main()
